#!/usr/bin/python
#coding=utf-8
'''
Studio della relazione
"volataggio applicato al piezo"-"posizione specchio"
nell'intervallo [v_min, v_max]:

---

Per ogni file nella cartella ./data/piezo viene eseguito il fit della figura 
di interferenza con la funzione
[0]+[1]*(cos(poln))^2

[0] = termine costante
[1] = termine di normalizzazione
poln = polinomiale di grado n (all'interno ha n+1 parametri) = f

---

Viene creato un albero con i parametri delle funzioni.

---

Vengono disegnati gli istogrammi dei valore dei parametri

'''

from __future__ import division, print_function
from ROOT import TCanvas, TGraphErrors, TF1, TLegend, TStyle, TTree, TFile, TH1D, TBrowser, TH2D, TStyle, gROOT
from math import pi
from array import array
from val2text import tex
from myanalisis import set_my_style
import os

'''
Definizione dei parametri e inizializzazione degli oggetti
'''
#Trasforma la fase del cos^2 nella distanza ottica introdotta 
l = 543 #nm
fase2dist = l/(2*pi)

#Formula per il fit della curva di interferenza
n = 2
v_min = 1.00
formula_prov = '[0]+[1]*(cos(pol1(2)*pol{0}(4)))^2'.format(n - 1)
formula_def = '[0]+[1]*(cos(pol{0}(2)))^2'.format(n)
functions = []
function_derivs = []

#Inizializza i file di risultati
path_graph = 'report/graph/piezo/'
path_table = 'report/table/piezo/'
path_value = 'report/value/piezo/'
path_results = 'results/piezo/'
ofile_name = path_results + 'piezo.txt'
ofile = open(ofile_name, 'w')
line = 'L\'argomento del coseno è una polinomiale di grado {0}\n'.format(n)
ofile.write(line)

#Inizializza l'albero dei parametri
fileroot = TFile(path_results + 'parameters.root','recreate')
tree_param  = TTree('par_parerr', 'Parametri e errori delle funzioni di risposta del piezo')
par_pol = array('i', [0])
parameters = array('d', (n + 1)*[0.])
sparameters = array('d', (n + 1)*[0.])
tree_param.Branch('n_pol', par_pol, 'n_pol/I')
tree_param.Branch('pars', parameters, 'pars[n_pol]/D')
tree_param.Branch('spars', sparameters, 'spars[n_pol]/D')
par_pol[0] = n + 1
par_list = []

'''
Fit su tutti i file che finiscono con .dat nella catella di dati
'''
#Iterazione sui file contenenti i dati
for root, dir, files in os.walk('./data/piezo'):
    for file in files:
        if file.endswith('.dat'):
            '''
            Fit del voltaggio generato dalla corrente del fotodiodo in funzione del
            voltaggio applicato al piezoelettrico.
            '''
            i_file = int(file.split('.')[1])
            ifile_name = os.path.join(root, file)
            graph_piezo = TGraphErrors(ifile_name)
            ofile.write('\nFit {0}\n'.format(i_file))

            #Fit provvisorio per ricavare i dati iniziali (i parametri non sono indipendenti)
            funz_piezo = TF1('funz_piezo', formula_prov, 0, 2)
            funz_piezo.SetParameters(200, 2300, 0, 4, 1, 0, 0)
            graph_piezo.SetTitle('interferenza {0};Ddp sul piezo [V];Ddp su R_L [mV]'.format(i_file))
            graph_piezo.Fit(funz_piezo, 'mq', '', v_min, 2)

            #Ricava i parametri iniziali per il fit con una polinomiale di grado n
            pars = [funz_piezo.GetParameter(i) for i in range(funz_piezo.GetNpar())]
            new_pars = []
            new_pars.append(pars[0])
            new_pars.append(pars[1])
            new_pars.append(pars[2]*pars[4])
            for i in range(1, n):
                par = pars[2]*pars[i+4] + pars[3]*pars[i+3]
                new_pars.append(par)
            new_pars.append(pars[3]*pars[3 + n])

            #Fit definitivo
            funz_finale = TF1('funz_finale', formula_def, 0.6, 2)
            for i in range(n+3):
                funz_finale.SetParameter(i, new_pars[i])

            graph_piezo.Fit(funz_finale, 'mq', '', v_min, 2)
            #canvas_piezo.SaveAs('risultati/inteferenza_{0}.pdf'.format(i_file))
            prob = funz_finale.GetProb()
            line = 'Probabilità del fit:        {0:.2}\n'.format(prob)     
            ofile.write(line)
            graph_piezo.GetXaxis().SetLimits(1, 2)
            if i_file == 15:
                canvas_piezo = TCanvas('canvas_piezo', 'canvas_piezo', 2)
                graph_piezo.Draw('ap')
                continue

            '''
            Studio della relazione Volt applicati al piezo-Distanza ottica introdotta
            '''
            funz_risposta = TF1('funz_risposta', '{1}*pol{0}'.format(n, fase2dist), 0, 2)
            ofile.write('\nFunzione di risposta {0}\nOrdine\tCoef\n'.format(i_file))
            for i in range(n + 1):
                parameter = funz_finale.GetParameter(i + 2)
                parameter_err = funz_finale.GetParError(i + 2)
                funz_risposta.SetParameter(i, parameter)
                line = '{0}\t{1:.3}\t±\t{2:.3}\n'.format(i, parameter, parameter_err)
                ofile.write(line)
                parameters[i] = parameter
                sparameters[i] = parameter_err

            tree_param.Fill()
            funz_risposta.SetTitle('Distanza ottica introdotta dal piezo {0};Ddp sul piezo [V];Posizione dello specchio[nm]'.format(i_file))
            funz_risposta.SetLineColor(i_file + 2)
            funz_risposta.SetParameter(0,0)
            functions.append(funz_risposta)
            par_list.append(parameters.tolist())


            '''
            Studio della derivata della relazione Volt applicati al piezo-Distanza ottica introdotta
            '''
            funz_deriv_risposta = TF1('funz_deriv_risposta', '{1}*pol{0}'.format(n - 1, fase2dist), 0, 2)
            ofile.write('\nDerivata della funzione di risposta {0}\nOrdine\tCoef\n'.format(i_file))
            for i in range(n):
                parameter = funz_risposta.GetParameter(i+1)
                funz_deriv_risposta.SetParameter(i, (i+1)*parameter)
                line = '{0}\t{1}\n'.format(i, (i+1)*parameter)
                ofile.write(line)

            funz_deriv_risposta.SetTitle('Derivata della distanza ottica introdotta dal piezo {0};Ddp sul piezo [V];[nm/V]'.format(i_file))
            funz_deriv_risposta.SetLineColor(i_file +2)
            function_derivs.append(funz_deriv_risposta)

fileroot.Write()
fileroot.Close()

'''
Analisi dei parametri delle funzioni di risposta
'''
#Calcola la media dei coefficienti
mean_pars = [0.]*(n + 1) 
for coefs in par_list:
    for i, coef in enumerate(coefs):
        mean_pars[i] += coef

mean_pars = [mean_par/len(par_list) for mean_par in mean_pars]

#Calcola la varianza dei coefficienti
variances = [0.]*(n + 1) 
for coefs in par_list:
    for i, coef in enumerate(coefs):
        variances[i] += (coef - mean_pars[i])**2

s_pars = [(variance/(len(par_list) - 1))**.5 for variance in variances]

#Scrive i valori dei parametri con gli errori nei file 
#Risultati
ofile_coef = open(path_results + 'par_spar.pol{0}.txt'.format(n), 'w')
for i in range(1, n + 1):
    par_spar = '{0}\t{1}\n'.format(mean_pars[i], s_pars[i])
    ofile_coef.write(par_spar)

ofile_coef.close()
#Relazione
for i in range(1, n + 1):
    tex(mean_pars[i], s_pars[i], path_value + 'funz_risposta_par{0}.txt'.format(i))
    
'''
Grafici
'''
#Creazione dello stile dei grafici
set_my_style()

#Disegna gli istogrammi con i parametri
fileroot = TFile(path_results + 'parameters.root')
tree_param  = fileroot.Get('par_parerr')
canvas_parametri = TCanvas('canvas_parametri', 'canvas_parametri', 2)
#Bidimensionale per del termine lineare e del suo errore
hist_par_and_err1 = TH2D('hist_par_and_err1',';Coef. lineare [rad/V];\sigma_{Coef. lineare} [rad/V]', 100, 3, 3.6, 100, 0.06 ,0.085)
tree_param.Project('hist_par_and_err1', 'spars[1]:pars[1]')
hist_par_and_err1.GetYaxis().SetTitleOffset(1.2)
hist_par_and_err1.Draw()
canvas_parametri.SaveAs(path_graph + 'par1_vs_spar1.pdf')
#Termine lineare
hist_par1 = TH1D('hist_par1',';Coef. lineare [rad/V];', 6, 2.9, 3.8)
tree_param.Project('hist_par1', 'pars[1]')
hist_par1.Draw()
canvas_parametri.SaveAs(path_graph + 'par1.pdf')
#Termine quadratico
hist_par2 = TH1D('hist_par2',';Coef. quadratico [rad/V^{2}];', 6, 0 ,0.3)
tree_param.Project('hist_par2', 'pars[2]')
hist_par2.Draw()
canvas_parametri.SaveAs(path_graph + 'par2.pdf')

#Disegna i fasci di funzioni
#Attenzione!!! I file non son processati in ordine alfabetico
canvas_func = TCanvas('canvas_func', 'canvas_func', 2)
funcset = functions
funcset[0].Draw()
for func in funcset:
    func.Draw('same')

legenda = TLegend(0.05, 0.85, 0.25, 0.7)
for i, func in enumerate(funcset):
    legenda.AddEntry(func, str(i + 1))

legenda.Draw()

#Termine del programma
ofile.close()
